#include "net/ip.h"
#include "net/udp.h"
#include "net/tcp.h"
#include "drivers/pcap_adapter.h"
#include "net/arp.h"
#include <string.h>
#include <stdlib.h>
#include <stdio.h>
#include <time.h> // added for random MAC generation

#ifdef _WIN32
/* do NOT include winsock2/ws2tcpip — we provide local byte-order helpers */
#include <windows.h>
#else
#include <sys/time.h>
#endif

#pragma pack(push,1)
typedef struct {
	uint8_t dst_mac[6];
	uint8_t src_mac[6];
	uint16_t ethertype;
	// IPv4 header follows
} eth_hdr_t;

typedef struct {
	uint8_t ver_ihl;
	uint8_t tos;
	uint16_t tot_len;
	uint16_t id;
	uint16_t frag_off;
	uint8_t ttl;
	uint8_t proto;
	uint16_t checksum;
	uint32_t saddr;
	uint32_t daddr;
	// options...
} ip_hdr_t;
#pragma pack(pop)

/* ----------------- local byte-order helpers -----------------
   Provide htons/ntohs/htonl/ntohl equivalents without system headers.
   Detect endianness at runtime and swap bytes on little-endian hosts.
*/
static int _is_little_endian(void) {
	unsigned int x = 1;
	return (*(unsigned char*)&x) == 1;
}

static uint16_t _swap16(uint16_t v) {
	return (uint16_t)((v << 8) | (v >> 8));
}

static uint32_t _swap32(uint32_t v) {
	return ((v << 24) & 0xFF000000UL) |
	       ((v << 8)  & 0x00FF0000UL) |
	       ((v >> 8)  & 0x0000FF00UL) |
	       ((v >> 24) & 0x000000FFUL);
}

static uint16_t my_htons(uint16_t v) {
	if (_is_little_endian()) return _swap16(v);
	return v;
}
static uint16_t my_ntohs(uint16_t v) { return my_htons(v); }

static uint32_t my_htonl(uint32_t v) {
	if (_is_little_endian()) return _swap32(v);
	return v;
}
static uint32_t my_ntohl(uint32_t v) { return my_htonl(v); }
/* ----------------------------------------------------------- */

static void* g_netif = NULL;
static uint32_t g_src_ip = 0;
static uint8_t g_src_mac[6];
static ip_payload_cb g_udp_cb = NULL;
static void* g_udp_ctx = NULL;
static void (*g_tcp_cb)(const uint8_t*, size_t, uint32_t, uint16_t, void*) = NULL;
static void* g_tcp_ctx = NULL;

/* 新增：强制目标 MAC 控制（默认不启用；demo 可通过接口启用并设置） */
static uint8_t g_forced_dst_mac[6] = { 0xFF,0xFF,0xFF,0xFF,0xFF,0xFF };
static int g_use_forced_dst_mac = 0;

static uint16_t ip_checksum(const void* vdata, size_t length) {
	// standard Internet checksum
	const uint8_t* data = (const uint8_t*)vdata;
	uint32_t sum = 0;
	while (length > 1) {
		sum += (data[0] << 8) | data[1];
		data += 2;
		length -= 2;
	}
	if (length > 0) {
		sum += data[0] << 8;
	}
	while (sum >> 16) sum = (sum & 0xFFFF) + (sum >> 16);
	return (uint16_t)(~sum);
}

int ip_init(void* netif_handle, uint32_t src_ip, const uint8_t src_mac[6]) {
	g_netif = netif_handle;
	g_src_ip = src_ip;
	/* 如果传入的 src_mac 全为 0 或为 NULL，则生成一个本地管理 (locally administered) 随机 MAC */
	int all_zero = 1;
	if (src_mac) {
		for (int i = 0; i < 6; ++i) {
			if (src_mac[i] != 0) { all_zero = 0; break; }
		}
	} else {
		all_zero = 1;
	}
	if (!src_mac || all_zero) {
		/* 生成 02:xx:xx:xx:xx:xx (locally administered) */
		unsigned int seed = (unsigned int)time(NULL) ^ (unsigned int)(uintptr_t)netif_handle;
		srand(seed);
		g_src_mac[0] = 0x02; /* locally administered, unicast */
		for (int i = 1; i < 6; ++i) g_src_mac[i] = rand() & 0xFF;
	} else {
		memcpy(g_src_mac, src_mac, 6);
	}
	udp_init();
	tcp_init();
	arp_init(g_netif, g_src_mac, g_src_ip); // 初始化 ARP
	return 0;
}

int ip_register_udp_cb(ip_payload_cb cb, void* ctx) {
	g_udp_cb = cb; g_udp_ctx = ctx; return 0;
}

int ip_register_tcp_cb(void (*cb)(const uint8_t*, size_t, uint32_t, uint16_t, void*), void* ctx) {
	g_tcp_cb = cb; g_tcp_ctx = ctx; return 0;
}

/*
 To bypass ARP and directly fill destination MAC, define SKIP_ARP_DIRECT at compile time.
 You can also override DEFAULT_DST_MAC by defining DEFAULT_DST_MAC_BYTES (as comma-separated hex bytes)
 when compiling, e.g. -DDEFAULT_DST_MAC_BYTES=0x00,0x11,0x22,0x33,0x44,0x55
*/
#ifndef SKIP_ARP_DIRECT
/* default: use ARP resolution path */
#else
/* default target MAC when SKIP_ARP_DIRECT is enabled */
#ifndef DEFAULT_DST_MAC_BYTES
#define DEFAULT_DST_MAC_BYTES 0xff,0xff,0xff,0xff,0xff,0xff
#endif
static const uint8_t DEFAULT_DST_MAC[6] = { DEFAULT_DST_MAC_BYTES };
#endif

int ip_send_udp(uint32_t dst_ip, uint16_t dst_port, uint16_t src_port, const uint8_t* payload, size_t len) {
	// Build Ethernet + IP + UDP into a frame, but DO NOT broadcast for unicast: use ARP resolution
	size_t ip_hdr_len = sizeof(ip_hdr_t);
	size_t udp_hdr_len = 8;
	size_t eth_len = sizeof(eth_hdr_t);
	size_t tot = eth_len + ip_hdr_len + udp_hdr_len + len;
	uint8_t* buf = (uint8_t*)malloc(tot);
	if (!buf) return -1;
	eth_hdr_t* eth = (eth_hdr_t*)buf;
	// dst mac to be filled by ARP if known; set to 00:00:00.. as placeholder
	memset(eth->dst_mac, 0, 6);
	memcpy(eth->src_mac, g_src_mac, 6);
	eth->ethertype = my_htons(0x0800);
	ip_hdr_t* iph = (ip_hdr_t*)(buf + eth_len);
	iph->ver_ihl = (4<<4) | (5);
	iph->tos = 0;
	iph->tot_len = my_htons((uint16_t)(ip_hdr_len + udp_hdr_len + len));
	iph->id = 0;
	iph->frag_off = 0;
	iph->ttl = 64;
	iph->proto = 17;
	iph->checksum = 0;
	iph->saddr = my_htonl(g_src_ip);
	iph->daddr = my_htonl(dst_ip);
	iph->checksum = ip_checksum(iph, ip_hdr_len);
	// UDP header
	uint8_t* udpp = (uint8_t*)iph + ip_hdr_len;
	uint16_t* up = (uint16_t*)udpp;
	up[0] = my_htons(src_port);
	up[1] = my_htons(dst_port);
	up[2] = my_htons((uint16_t)(udp_hdr_len + len));
	up[3] = 0; // checksum optional/zero
	memcpy(udpp + udp_hdr_len, payload, len);

#ifdef SKIP_ARP_DIRECT
	/* Directly set destination MAC and send (skip ARP) */
	memcpy(eth->dst_mac, DEFAULT_DST_MAC, 6);
	int r = pcap_adapter_send(g_netif, buf, tot);
	free(buf);
	return r;
#else
	/* 如果启用了强制目标 MAC，优先直接发送 */
	if (g_use_forced_dst_mac) {
		memcpy(eth->dst_mac, g_forced_dst_mac, 6);
		int r = pcap_adapter_send(g_netif, buf, tot);
		free(buf);
		return r;
	}
	/* try to resolve MAC via ARP cache */
	uint8_t dst_mac[6];
	if (arp_resolve(dst_ip, dst_mac) == 0) {
		memcpy(eth->dst_mac, dst_mac, 6);
		int r = pcap_adapter_send(g_netif, buf, tot);
		free(buf);
		return r;
	} else {
		/* queue frame and trigger ARP request */
		int r = arp_queue_frame(dst_ip, buf, tot);
		if (r != 0) free(buf);
		return r;
	}
#endif
}

int ip_input(const uint8_t* pkt, size_t len) {
	if (!pkt || len < sizeof(eth_hdr_t) + sizeof(ip_hdr_t)) return -1;
	const eth_hdr_t* eth = (const eth_hdr_t*)pkt;
	uint16_t ethertype = my_ntohs(eth->ethertype);
	if (ethertype == 0x0806) { // ARP
		arp_input(pkt, len);
		return 0;
	}
	if (ethertype != 0x0800) return 0; // not IPv4
	const ip_hdr_t* iph = (const ip_hdr_t*)(pkt + sizeof(eth_hdr_t));
	int ihl = (iph->ver_ihl & 0x0F) * 4;
	if (len < sizeof(eth_hdr_t) + ihl) return -1;
	uint32_t saddr = my_ntohl(iph->saddr);
	uint32_t daddr = my_ntohl(iph->daddr);
	if (iph->proto == 17) {
		// UDP
		const uint8_t* udpp = (const uint8_t*)iph + ihl;
		if ((size_t)((len) - sizeof(eth_hdr_t) - ihl) < 8) return -1;
		uint16_t src_port = my_ntohs(*(uint16_t*)(udpp));
		uint16_t dst_port = my_ntohs(*(uint16_t*)(udpp+2));
		uint16_t ulen = my_ntohs(*(uint16_t*)(udpp+4));
		const uint8_t* payload = udpp + 8;
		size_t plen = ulen - 8;
		if (g_udp_cb) g_udp_cb(payload, plen, saddr, src_port, g_udp_ctx);
	} else if (iph->proto == 6) {
		// TCP
		const uint8_t* tcpp = (const uint8_t*)iph + ihl;
		size_t tlen = len - sizeof(eth_hdr_t) - ihl;
		if (g_tcp_cb) g_tcp_cb(tcpp, tlen, saddr, 0, g_tcp_ctx);
	}
	return 0;
}

// 新增：IP层接收来自 TCP 层要发送的 TCP 段（IP+以太封装并通过ARP发送或队列）
int ip_send_tcp(uint32_t dst_ip, uint16_t dst_port, uint16_t src_port, const uint8_t* tcp_segment, size_t seg_len) {
	size_t ip_hdr_len = 20;
	size_t eth_len = sizeof(eth_hdr_t);
	size_t tot = eth_len + ip_hdr_len + seg_len;
	uint8_t* buf = (uint8_t*)malloc(tot);
	if (!buf) return -1;
	// eth dst filled later by ARP; set placeholder dst=00:00:00:00:00:00
	memset(buf, 0, 6);
	memcpy(buf+6, g_src_mac, 6);
	*(uint16_t*)(buf+12) = my_htons(0x0800);
	// ip hdr
	uint8_t* iph = buf + eth_len;
	iph[0] = (4<<4) | 5;
	iph[1] = 0;
	uint16_t totallen = my_htons((uint16_t)(ip_hdr_len + seg_len));
	memcpy(iph+2, &totallen, 2);
	*(uint16_t*)(iph+4) = 0;
	iph[8] = 64;
	iph[9] = 6;
	*((uint32_t*)(iph+12)) = my_htonl(g_src_ip);
	*((uint32_t*)(iph+16)) = my_htonl(dst_ip);
	// copy tcp segment (assumed contains tcp header+payload)
	memcpy(buf + eth_len + ip_hdr_len, tcp_segment, seg_len);

#ifdef SKIP_ARP_DIRECT
	/* Fill destination MAC directly and send */
	memcpy(buf, DEFAULT_DST_MAC, 6);
	int r = pcap_adapter_send(g_netif, buf, tot);
	free(buf);
	return r;
#else
	/* 如果启用了强制目标 MAC，优先直接发送 */
	if (g_use_forced_dst_mac) {
		memcpy(buf, g_forced_dst_mac, 6);
		int r = pcap_adapter_send(g_netif, buf, tot);
		free(buf);
		return r;
	}
	/* try to resolve MAC */
	uint8_t dst_mac[6];
	if (arp_resolve(dst_ip, dst_mac) == 0) {
		memcpy(buf, dst_mac, 6);
		int r = pcap_adapter_send(g_netif, buf, tot);
		free(buf);
		return r;
	} else {
		int r = arp_queue_frame(dst_ip, buf, tot);
		if (r != 0) free(buf);
		return r;
	}
#endif
}

/* 新增：设置/启用强制目标 MAC 的接口（demo 调用） */
int ip_set_forced_dst_mac(const uint8_t mac[6]) {
	if (!mac) return -1;
	memcpy(g_forced_dst_mac, mac, 6);
	return 0;
}

int ip_set_use_forced_dst_mac(int use) {
	g_use_forced_dst_mac = use ? 1 : 0;
	return 0;
}
